#!/usr/bin/env python3

import os
from ament_index_python.packages import get_package_share_directory
from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument, OpaqueFunction
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution
from launch_ros.actions import Node
from launch.conditions import IfCondition

def launch_setup(context, *args, **kwargs):
    """Setup function for launch"""
    
    # Get package directories
    ego_adapter_dir = get_package_share_directory('ego_planner_adapter')
    nav_bringup_dir = get_package_share_directory('pb2025_nav_bringup')
    
    # Parameters
    use_sim_time = LaunchConfiguration('use_sim_time')
    params_file = LaunchConfiguration('params_file')
    namespace = LaunchConfiguration('namespace')
    
    # Configuration files
    trajectory_tracker_params = PathJoinSubstitution([
        ego_adapter_dir, 'config', 'trajectory_tracker_params.yaml'
    ])
    
    target_relay_params = PathJoinSubstitution([
        ego_adapter_dir, 'config', 'target_relay_params.yaml'
    ])
    
    # Ego-planner nodes
    ego_planner_node = Node(
        package='ego_planner',
        executable='ego_planner_node',
        name='ego_planner',
        namespace=namespace,
        parameters=[
            {'use_sim_time': use_sim_time},
            params_file
        ],
        remappings=[
            ('odometry', 'odometry'),
            ('cloud', 'cloud_registered'),
            ('waypoint_generator/waypoints', 'planning/waypoints'),
            ('planning/bspline', 'planning/trajectory')
        ],
        output='screen'
    )
    
    # Trajectory tracker node
    trajectory_tracker_node = Node(
        package='ego_planner_adapter',
        executable='trajectory_tracker_node',
        name='trajectory_tracker',
        namespace=namespace,
        parameters=[
            {'use_sim_time': use_sim_time},
            trajectory_tracker_params
        ],
        remappings=[
            ('odometry', 'odometry'),
            ('trajectory', 'planning/trajectory'),
            ('cmd_vel', 'cmd_vel_nav2_result'),
            ('target_altitude', 'target_altitude')
        ],
        output='screen'
    )
    
    # Target relay node
    target_relay_node = Node(
        package='ego_planner_adapter',
        executable='target_relay_node',
        name='target_relay',
        namespace=namespace,
        parameters=[
            {'use_sim_time': use_sim_time},
            target_relay_params
        ],
        remappings=[
            ('goal_pose', 'goal_pose'),
            ('target_altitude', 'target_altitude'),
            ('waypoint_generator/waypoints', 'planning/waypoints'),
            ('planning_trigger', 'planning/trigger')
        ],
        output='screen'
    )
    
    # Grid map node (for ego-planner)
    grid_map_node = Node(
        package='plan_env',
        executable='obj_generator',
        name='obj_generator',
        namespace=namespace,
        parameters=[
            {'use_sim_time': use_sim_time},
            params_file
        ],
        remappings=[
            ('odometry', 'odometry'),
            ('cloud', 'cloud_registered'),
            ('grid_map/odom', 'odometry'),
            ('grid_map/cloud', 'cloud_registered')
        ],
        output='screen'
    )
    
    # Keep existing essential nodes
    livox_driver_node = Node(
        package='livox_ros_driver2',
        executable='livox_ros_driver2_node',
        name='livox_ros_driver2_node',
        parameters=[
            nav_bringup_dir + '/config/reality/nav2_params.yaml'
        ],
        output='screen'
    )
    
    point_lio_node = Node(
        package='point_lio',
        executable='pointlio_mapping',
        name='pointlio_mapping',
        parameters=[
            nav_bringup_dir + '/config/reality/nav2_params.yaml'
        ],
        remappings=[
            ('aft_mapped_to_init', 'odometry')
        ],
        output='screen'
    )
    
    loam_interface_node = Node(
        package='loam_interface',
        executable='loam_interface',
        name='loam_interface',
        parameters=[
            nav_bringup_dir + '/config/reality/nav2_params.yaml'
        ],
        output='screen'
    )
    
    # Keep drone serial interface
    drone_serial_node = Node(
        package='drone_serial_interface',
        executable='drone_serial_interface_node',
        name='drone_serial_interface_node',
        parameters=[
            nav_bringup_dir + '/config/reality/nav2_params.yaml'
        ],
        remappings=[
            ('cmd_vel', 'cmd_vel_nav2_result')
        ],
        output='screen'
    )
    
    # Fake velocity transform
    fake_vel_transform_node = Node(
        package='fake_vel_transform',
        executable='fake_vel_transform',
        name='fake_vel_transform',
        parameters=[
            nav_bringup_dir + '/config/reality/nav2_params.yaml'
        ],
        output='screen'
    )
    
    # Return all nodes
    return [
        livox_driver_node,
        point_lio_node,
        loam_interface_node,
        grid_map_node,
        ego_planner_node,
        trajectory_tracker_node,
        target_relay_node,
        drone_serial_node,
        fake_vel_transform_node
    ]

def generate_launch_description():
    """Generate launch description"""
    
    # Declare launch arguments
    declare_use_sim_time = DeclareLaunchArgument(
        'use_sim_time',
        default_value='false',
        description='Use simulation time'
    )
    
    declare_params_file = DeclareLaunchArgument(
        'params_file',
        default_value=PathJoinSubstitution([
            get_package_share_directory('pb2025_nav_bringup'),
            'config', 'reality', 'nav2_params.yaml'
        ]),
        description='Full path to the ROS2 parameters file'
    )
    
    declare_namespace = DeclareLaunchArgument(
        'namespace',
        default_value='',
        description='Top-level namespace'
    )
    
    return LaunchDescription([
        declare_use_sim_time,
        declare_params_file,
        declare_namespace,
        OpaqueFunction(function=launch_setup)
    ])